from numpy import *
from Radonpy import *
import matplotlib.pyplot as plt
from math import pi
def hann(inp,prc):
	L=inp.size
	lw=ceil(L*prc/200.0)
	up=L-1-lw
	hn=ones(L)
	t=arange(0,lw)
	hn[0:lw]=hn[0:lw]-sin(pi/2*(lw-t)/lw)**2
	hn[L-1:up:-1]=hn[L-1:up:-1]-sin(pi/2*(lw-t)/lw)**2
	out=inp*hn
	return out
def matchfilspec(src,target,un,ln,dmp):
	srcf=fft.rfft(src)
	targf=fft.rfft(target)
	mp=max(abs(srcf))**2
	noise=ones(srcf.shape)*ln*mp
	w=arange(0,srcf.shape)
	noise=noise+un*mp*exp(-float(w)/dmp)
	ff=srcf.conj()*targf/(srcf.conj()*srcf+noise)
	mf=fft.irfft(ff)
	return mf
	
def matchfil(src,target,len,un,ln,dmp):
	L=src.size
	if L!=target.size:
		print "The size of input and that of target don't match!"
		return 0
	if len>L:
		print "The filter lenght is larger than that of input wavelet"
		return 0
	srcf=fft.rfft(src)
	targf=fft.rfft(target)
	targf=fft.rfft(target)
	mp=max(abs(srcf))**2
	noise=ones(srcf.shape)*ln*mp
	w=arange(0,srcf.size)
	noise=noise+un*mp*exp(-w/float(dmp))
	nm=fft.irfft(noise)
	srcautoc=fft.irfft(srcf.conj()*srcf)
	rw=srcautoc+nm
	stcorr=fft.irfft(srcf.conj()*targf)
	stcorr=roll(stcorr,L/2)
	mss=0
	strt=0
	for i in range(0,L-len+1):
		ssv=sum(abs(stcorr[i:len+i]))
		if ssv>mss:
			mss=ssv
			strt=i
	y=stcorr[i:i+len]
	A=zeros((len,len))
	for i in range(0,len):
		for j in range(0,len):
			k=abs(i-j)
			A[i,j]=rw[k]
	f=linalg.solve(A,y)
	zl=L/2-k
	return (f,zl)

if __name__=="__main__":
	L=256
	t=arange(0,L)
	target=rickerwav(t,128,3)+1e-10*random.randn(L)
	src=phaserot(roll(target,3),0.2)+0.05*gausswav(t,150,3)+0.02*gausswav(t,160,3)
	tf=fft.rfft(target)
	sf=fft.rfft(src)
#	plt.plot(log10(abs(tf)))
#	plt.plot(log10(abs(sf)))
#	plt.show()
	fil,zp=matchfil(src,target,256,0.1,1e-16,3)
	filhn=hann(fil,20)
	print filhn
#	plt.plot(fil)
#	plt.plot(filhn)
#	print zp
	plt.plot(log10(abs(fft.rfft(fil))))
	plt.plot(log10(abs(fft.rfft(filhn))))
#	plt.show()
	filsrc=convolve(fil,src)
#	plt.plot(log10(abs(fft.rfft(filsrc))))
	plt.show()
